﻿#define _CRT_SECURE_NO_WARNINGS
#include <winsock2.h>
#include <ws2tcpip.h>
#include <windows.h>
#include <iostream>
#include <map>
#include <set>
#include <string>
#include <ctime>
#include <vector> // 包含 vector 用于清理
#pragma comment(lib, "ws2_32.lib")

#define PORT 8888
#define HEARTBEAT_INTERVAL 10 // 为了测试，设置为10秒
#define MAX_ALLOWED_DEVICES 2 // 为了测试，设置为2

struct LicenseInfo {
    int max_users = MAX_ALLOWED_DEVICES; // 使用宏初始化
    std::set<std::string> active_users;
    std::map<std::string, time_t> last_heartbeat;
};

std::map<std::string, LicenseInfo> licenses; // 全局存储许可证信息

// 生成10位数字序列号
std::string generate_serial() {
    srand((unsigned)time(0));
    std::string serial;
    for (int i = 0; i < 10; ++i) {
        serial += '0' + (rand() % 10);
    }
    return serial;
}

// 将许可证信息保存到文件
void save_licenses() {
    FILE* f = fopen("licenses.dat", "w");
    if (!f) {
        perror("Error opening licenses.dat for writing");
        return;
    }
    // 使用传统迭代器循环代替 C++17 结构化绑定
    for (std::map<std::string, LicenseInfo>::const_iterator it = licenses.begin(); it != licenses.end(); ++it) {
        const std::string& serial = it->first;
        const LicenseInfo& info = it->second;
        fprintf(f, "%s %d\n", serial.c_str(), info.max_users);
        // 同样使用传统迭代器
        for (std::set<std::string>::const_iterator user_it = info.active_users.begin(); user_it != info.active_users.end(); ++user_it) {
            const std::string& user = *user_it;
            // 确保 user 存在于 last_heartbeat 中
            std::map<std::string, time_t>::const_iterator hb_it = info.last_heartbeat.find(user);
            if (hb_it != info.last_heartbeat.end()) {
                // 使用 %lld 匹配 time_t (__int64) 并进行类型转换
                fprintf(f, "U %s %lld\n", user.c_str(), (long long)hb_it->second);
            }
        }
    }
    fclose(f);
}

// 从文件加载许可证信息
void load_licenses() {
    FILE* f = fopen("licenses.dat", "r");
    if (!f) {
        return; // 文件不存在是正常情况
    }
    char serial_buf[32] = { 0 }; // 使用不同的名字避免歧义
    int max_users_from_file;
    // 读取主许可证信息
    while (fscanf(f, "%31s %d", serial_buf, &max_users_from_file) == 2) {
        std::string current_serial(serial_buf); // 转换为 std::string
        LicenseInfo info;
        info.max_users = max_users_from_file;
        licenses[current_serial] = info; // 先插入

        // 读取关联的用户信息
        char line_type[2] = { 0 };
        char username_buf[32] = { 0 }; // 使用不同的名字
        long long t = 0;
        long pos = ftell(f); // 保存当前文件位置

        // 尝试读取下一行判断是否为 'U'
        if (fscanf(f, "%1s", line_type) == 1 && line_type[0] == 'U') {
            // 如果是 'U'，则读取用户名和时间戳
            if (fscanf(f, "%31s %lld", username_buf, &t) == 2) {
                // 显式创建 std::string 插入 set
                licenses[current_serial].active_users.insert(std::string(username_buf));
                licenses[current_serial].last_heartbeat[std::string(username_buf)] = (time_t)t;
            }
            else {
                // 读取失败，回退文件指针
                fseek(f, pos, SEEK_SET);
            }
        }
        else {
            // 不是 'U' 开头，回退文件指针
            fseek(f, pos, SEEK_SET);
        }
        // 清空缓冲区，准备读取下一个序列号
        memset(serial_buf, 0, sizeof(serial_buf));
    }
    fclose(f);
}


// 清理超时的客户端连接
void cleanup_expired() {
    time_t now = time(0);
    bool changed = false;
    // 使用传统迭代器
    for (std::map<std::string, LicenseInfo>::iterator it = licenses.begin(); it != licenses.end(); ++it) {
        const std::string& serial = it->first;
        LicenseInfo& info = it->second; // 需要非 const 引用来修改
        std::vector<std::string> to_remove; // 使用 vector 存储待删除用户
        // 使用传统迭代器
        for (std::set<std::string>::const_iterator user_it = info.active_users.begin(); user_it != info.active_users.end(); ++user_it) {
            const std::string& user = *user_it;
            std::map<std::string, time_t>::const_iterator hb_it = info.last_heartbeat.find(user);
            if (hb_it != info.last_heartbeat.end()) {
                if (now - hb_it->second > HEARTBEAT_INTERVAL) {
                    to_remove.push_back(user); // 添加到待删除列表
                    std::cout << "User " << user << " on serial " << serial << " timed out." << std::endl;
                }
            }
            else {
                to_remove.push_back(user); // 没有心跳记录也移除
                std::cout << "User " << user << " on serial " << serial << " has no heartbeat record, removing." << std::endl;
            }
        }
        // 执行删除
        for (size_t i = 0; i < to_remove.size(); ++i) {
            info.active_users.erase(to_remove[i]);
            info.last_heartbeat.erase(to_remove[i]);
            changed = true;
        }
    }
    if (changed) {
        save_licenses(); // 如果有用户被移除，更新文件
    }
}

// 处理客户端请求的线程函数
DWORD WINAPI client_thread(LPVOID lpParam) {
    SOCKET client = (SOCKET)lpParam;
    char buf[256];
    int len = recv(client, buf, sizeof(buf) - 1, 0);

    if (len <= 0) {
        closesocket(client);
        return 1;
    }
    buf[len] = '\0';
    std::string cmd_str(buf); // 使用 std::string 处理命令
    std::cout << "Received command: " << cmd_str << std::endl;

    char command_type[32] = { 0 };
    char serial_buf[32] = { 0 }; // 缓冲区变量名区分
    char username_buf[32] = { 0 };
    char password_buf[32] = { 0 };
    char type_buf[16] = { 0 };

    // 解析命令类型
    sscanf(buf, "%31s", command_type);

    if (strcmp(command_type, "REGISTER") == 0) {
        if (sscanf(buf, "%*s %31s %31s %15s", username_buf, password_buf, type_buf) == 3) {
            std::string new_serial = generate_serial();
            LicenseInfo info;
            licenses[new_serial] = info;
            save_licenses();
            // 使用 (int) 转换 size_t
            send(client, new_serial.c_str(), (int)new_serial.size(), 0);
            std::cout << "Registered user " << username_buf << " with serial " << new_serial << std::endl;
        }
        else {
            send(client, "ERROR Invalid REGISTER format", 28, 0);
        }
    }
    else if (strcmp(command_type, "AUTH") == 0) {
        if (sscanf(buf, "%*s %31s %31s", serial_buf, username_buf) == 2) {
            cleanup_expired();
            std::string serial_str(serial_buf); // 转为 string
            std::string username_str(username_buf); // 转为 string
            if (licenses.count(serial_str) == 0) {
                send(client, "DENY Invalid Serial", 19, 0);
                std::cout << "Auth failed for " << username_str << ": Invalid serial " << serial_str << std::endl;
            }
            else {
                LicenseInfo& info = licenses[serial_str]; // 需要非 const 引用
                // 使用 size() 获取 set 大小
                if (info.active_users.size() < info.max_users || info.active_users.count(username_str)) {
                    info.active_users.insert(username_str); // 插入 string
                    info.last_heartbeat[username_str] = time(0);
                    save_licenses();
                    send(client, "ALLOW", 5, 0);
                    std::cout << "Auth success for " << username_str << " on serial " << serial_str << std::endl;
                }
                else {
                    send(client, "DENY Max Users Reached", 22, 0);
                    std::cout << "Auth failed for " << username_str << ": Max users reached on serial " << serial_str << std::endl;
                }
            }
        }
        else {
            send(client, "ERROR Invalid AUTH format", 24, 0);
        }
    }
    else if (strcmp(command_type, "HEARTBEAT") == 0) {
        if (sscanf(buf, "%*s %31s %31s", serial_buf, username_buf) == 2) {
            std::string serial_str(serial_buf);
            std::string username_str(username_buf);
            if (licenses.count(serial_str) && licenses[serial_str].active_users.count(username_str)) {
                licenses[serial_str].last_heartbeat[username_str] = time(0);
                send(client, "OK", 2, 0);
                std::cout << "Heartbeat received from " << username_str << " on serial " << serial_str << std::endl;
            }
            else {
                send(client, "ERROR Invalid Heartbeat", 23, 0);
                std::cout << "Invalid heartbeat from " << username_str << " on serial " << serial_str << std::endl;
            }
        }
        else {
            send(client, "ERROR Invalid HEARTBEAT format", 29, 0);
        }
    }
    else if (strcmp(command_type, "LOGOUT") == 0) {
        if (sscanf(buf, "%*s %31s %31s", serial_buf, username_buf) == 2) {
            std::string serial_str(serial_buf);
            std::string username_str(username_buf);
            if (licenses.count(serial_str)) {
                licenses[serial_str].active_users.erase(username_str);
                licenses[serial_str].last_heartbeat.erase(username_str);
                save_licenses();
                send(client, "OK", 2, 0);
                std::cout << "User " << username_str << " logged out from serial " << serial_str << std::endl;
            }
            else {
                send(client, "OK", 2, 0);
            }
        }
        else {
            send(client, "ERROR Invalid LOGOUT format", 26, 0);
        }
    }
    else {
        send(client, "ERROR Unknown Command", 21, 0);
        std::cout << "Received unknown command: " << cmd_str << std::endl;
    }

    closesocket(client);
    return 0;
}

int main() {
    WSADATA wsa;
    // 检查 WSAStartup 返回值
    if (WSAStartup(MAKEWORD(2, 2), &wsa) != 0) {
        std::cerr << "WSAStartup failed. Error Code : " << WSAGetLastError() << std::endl;
        return 1;
    }

    load_licenses();

    SOCKET s = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (s == INVALID_SOCKET) {
        std::cerr << "Could not create socket : " << WSAGetLastError() << std::endl;
        WSACleanup();
        return 1;
    }

    sockaddr_in server_addr;
    server_addr.sin_family = AF_INET;
    server_addr.sin_addr.s_addr = INADDR_ANY;
    server_addr.sin_port = htons(PORT);

    if (bind(s, (sockaddr*)&server_addr, sizeof(server_addr)) == SOCKET_ERROR) {
        std::cerr << "Bind failed with error code : " << WSAGetLastError() << std::endl;
        closesocket(s);
        WSACleanup();
        return 1;
    }

    listen(s, SOMAXCONN);

    std::cout << "License server started on port " << PORT << std::endl;
    std::cout << "Max allowed devices per license: " << MAX_ALLOWED_DEVICES << std::endl;
    std::cout << "Heartbeat interval (seconds): " << HEARTBEAT_INTERVAL << std::endl;


    SOCKET client_socket;
    sockaddr_in client_addr;
    int client_addr_size = sizeof(client_addr);

    while (true) {
        client_socket = accept(s, (sockaddr*)&client_addr, &client_addr_size);
        if (client_socket == INVALID_SOCKET) {
            std::cerr << "accept failed with error code : " << WSAGetLastError() << std::endl;
            continue;
        }

        char client_ip[INET_ADDRSTRLEN];
        // 使用 inet_ntop 输出 IP 地址
        if (inet_ntop(AF_INET, &client_addr.sin_addr, client_ip, INET_ADDRSTRLEN) != NULL) {
            std::cout << "Connection accepted from " << client_ip << ":" << ntohs(client_addr.sin_port) << std::endl;
        }
        else {
            std::cout << "Connection accepted from unknown IP" << std::endl;
        }


        HANDLE hThread = CreateThread(NULL, 0, client_thread, (LPVOID)client_socket, 0, NULL);
        if (hThread == NULL) {
            std::cerr << "CreateThread failed for client. Error Code : " << GetLastError() << std::endl;
            closesocket(client_socket);
        }
        else {
            CloseHandle(hThread);
        }
    }

    closesocket(s);
    WSACleanup();
    return 0;
}